//
// Created by f on 2022/3/25.
//

#include "quadtree_map.h"

QuadtreeNode::QuadtreeNode(int min_x, int min_y, int size):
min_x_(min_x), min_y_(min_y), size_(size)
{
    lu_ = nullptr;
    ru_ = nullptr;
    ld_ = nullptr;
    rd_ = nullptr;
}

QuadtreeNode::QuadtreeNode(std::weak_ptr<QuadtreeNode> father, int min_x, int min_y, int size):
father_(father), min_x_(min_x), min_y_(min_y), size_(size)
{
    lu_ = nullptr;
    ru_ = nullptr;
    ld_ = nullptr;
    rd_ = nullptr;
}

void QuadtreeNode::SetLeaf(){
    value_ = 1;
    lu_.reset();
    ru_.reset();
    ld_.reset();
    rd_.reset();
}

bool QuadtreeNode::isLeaf(){
    return (value_ == 1 && lu_ == nullptr && ru_ == nullptr && ld_ == nullptr && rd_ == nullptr);
}

const int& QuadtreeNode::Min_x(){
    return std::cref(min_x_);
}

const int& QuadtreeNode::Min_y(){
    return std::cref(min_y_);
}

const int& QuadtreeNode::Size(){
    return std::cref(size_);
}

QuadtreeMap::QuadtreeMap(const double &center_x, const double &center_y, const double &center_yaw,
                         const double &resolution, const double &min_dim,
                         const int &layer_id):
                         center_x_(center_x), center_y_(center_y), center_yaw_(center_yaw),
                         grid_dim_(resolution), layer_id_(layer_id), MapBase()
{
    int min_size = min_dim/grid_dim_;
    while (map_size_ < min_size)
    {
        map_size_ *= 2;
        max_depth_ += 1;
    }
    root_ = std::make_shared<QuadtreeNode>(0, 0, map_size_);
}

QuadtreeMap::QuadtreeMap(const int &layer_id):layer_id_(layer_id), MapBase(){}

bool QuadtreeMap::AddPoint(const double& x, const double& y){
    int row=0, col=0;
    {
        ///转换到网格坐标系下
        double x_trans = x - center_x_;
        double y_trans = y - center_y_;
        double rot_angle = -center_yaw_;
        double x_rot = cos(rot_angle) * x_trans - sin(rot_angle) * y_trans;
        double y_rot = sin(rot_angle) * x_trans + cos(rot_angle) * y_trans;
        ///找到对应的网格index
        row = floor(x_rot / grid_dim_ + map_size_ / 2);
        col = floor(y_rot / grid_dim_ + map_size_ / 2);
    }

    if(row<0 || col<0 || row>map_size_-1 || col>map_size_-1){
//        printf("Point out of map range\n");
        return false;
    }

    std::shared_ptr< QuadtreeNode > curr_node = root_;
    int curr_depth = 0;
    ///不断深入，直到找到对应的叶子节点
    while(curr_depth < max_depth_){
        if( curr_node->isLeaf() ) return true; ///该网格已经被添加过点

        ///right up
        int center_row = curr_node->Min_x()+curr_node->Size()/2;
        int center_col = curr_node->Min_y()+curr_node->Size()/2;
        if(row >= center_row && col >= center_col){
            if(curr_node->ru_ == nullptr)
                curr_node->ru_ = std::make_shared< QuadtreeNode >
                        (curr_node, center_row, center_col,
                         curr_node->Size()/2);
                curr_node = curr_node->ru_;
                curr_depth += 1;
        }
        ///left up
        else if(row >= center_row && col < center_col){
            if(curr_node->lu_ == nullptr)
                curr_node->lu_ = std::make_shared< QuadtreeNode >
                        (curr_node, center_row, curr_node->Min_y(),
                         curr_node->Size()/2);
            curr_node = curr_node->lu_;
            curr_depth += 1;
        }
        ///right down
        else if(row < center_row && col >= center_col){
            if(curr_node->rd_ == nullptr)
                curr_node->rd_ = std::make_shared< QuadtreeNode >
                        (curr_node, curr_node->Min_x(), center_col,
                         curr_node->Size()/2);
            curr_node = curr_node->rd_;
            curr_depth += 1;
        }
        ///left down
        else{
            if(curr_node->ld_ == nullptr)
                curr_node->ld_ = std::make_shared< QuadtreeNode >
                        (curr_node, curr_node->Min_x(), curr_node->Min_y(),
                         curr_node->Size()/2);
            curr_node = curr_node->ld_;
            curr_depth += 1;
        }
    }
    ///此时curr_node已经指向对应的叶子节点
    curr_node->SetLeaf();
    ///自底向上遍历，合并值相同的节点
    while (curr_node->father_.lock() != nullptr){
        ///向上走一层
        curr_node = curr_node->father_.lock();
        if(curr_node->ru_!=nullptr && curr_node->lu_!=nullptr && curr_node->rd_!=nullptr && curr_node->ld_!=nullptr){
            if(curr_node->ru_->isLeaf() && curr_node->lu_->isLeaf() && curr_node->rd_->isLeaf() && curr_node->ld_->isLeaf()){
                curr_node->SetLeaf();
//                printf("merge node x:%d y:%d size:%d\n",curr_node->Min_x(), curr_node->Min_y(), curr_node->Size());
            }
        }
        else break;
    }
    return true;
}

uint8_t QuadtreeMap::GetValue(const double& x, const double& y){
    int row, col;
    {
        ///转换到网格坐标系下
        double x_trans = x - center_x_;
        double y_trans = y - center_y_;
        double rot_angle = -center_yaw_;
        double x_rot = cos(rot_angle) * x_trans - sin(rot_angle) * y_trans;
        double y_rot = sin(rot_angle) * x_trans + cos(rot_angle) * y_trans;
        ///找到对应的网格index
        row = floor(x_rot / grid_dim_ + map_size_ / 2);
        col = floor(y_rot / grid_dim_ + map_size_ / 2);
    }

    if(row<0 || col<0 || row>map_size_-1 || col>map_size_-1){
        printf("Point out of map range\n");
        return false;
    }

    std::shared_ptr< QuadtreeNode > curr_node = root_;
    int curr_depth = 0;
    uint8_t value;
    ///寻找对应的最小节点。只有两种可能：最小节点是空节点，返回0；最小节点是Leaf，返回1。
        while(curr_depth <= max_depth_){
        if( curr_node->isLeaf() ) {value=1; break;}

        int center_row = curr_node->Min_x()+curr_node->Size()/2;
        int center_col = curr_node->Min_y()+curr_node->Size()/2;
        ///right up
        if(row >= center_row && col >= center_col){
            if(curr_node->ru_ == nullptr) {value=0; break;}
            curr_node = curr_node->ru_;
            curr_depth += 1;
        }
        ///left up
        else if(row >= center_row && col < center_col){
            if(curr_node->lu_ == nullptr) {value=0; break;}
            curr_node = curr_node->lu_;
            curr_depth += 1;
        }
        ///right down
        else if(row < center_row && col >= center_col){
            if(curr_node->rd_ == nullptr) {value=0; break;}
            curr_node = curr_node->rd_;
            curr_depth += 1;
        }
        ///left down
        else{
            if(curr_node->ld_ == nullptr) {value=0; break;}
            curr_node = curr_node->ld_;
            curr_depth += 1;
        }
    }
    return value;
}

bool QuadtreeMap::SaveMap(const std::string& save_path,  const std::string& save_type){
    std::string map_file_path = save_path + "quadtreemap_"
                                + std::to_string(center_x_) + "_"
                                + std::to_string(center_y_) + "_"
                                + std::to_string(center_yaw_) + "_"
                                + std::to_string(grid_dim_);
    if(save_type == "PNG"){
        printf("Saving map file ... ...\n");
        map_file_path += ".png";
        cv::Mat m;
        toCVMat(m);
        cv::imwrite(map_file_path , m);
        printf("Map saved at %s\n", map_file_path.c_str());
        return true;
    }
    else if(save_type == "PGM"){
        printf("Saving map file ... ...\n");
        map_file_path += ".pgm";
        std::vector<int> compression_params;
        compression_params.emplace_back(CV_IMWRITE_PXM_BINARY); //Set to PXM compression
        compression_params.emplace_back(0); //Set type of PXM in our case PGM
        cv::Mat m;
        toCVMat(m);
        cv::imwrite(map_file_path, m, compression_params);
        printf("Map saved at %s\n", map_file_path.c_str());
        return true;
    }
    else if(save_type == "BIN"){
        printf("Saving map file ... ...\n");
        map_file_path += ".bin";

        std::ofstream map_file;
        map_file.open(map_file_path, std::ios_base::out|std::ios::binary);

        ///保存参数map_size_
        map_file.write((char *)&(map_size_), sizeof(int));

        ///层序遍历，找到所有叶子节点
        std::queue< std::shared_ptr<QuadtreeNode> > q;
        std::vector<int> leaf_node_params;
        int leaf_node_num = 0;
        q.push(root_);
        while (!q.empty()){
            int curr_level_size = q.size();
            for(int i=0; i<curr_level_size; i++){
                auto node = q.front(); q.pop();

                ///要先保存叶子节点的数量，所以这里先把叶子节点放在vector中
                if( node->isLeaf() ){
                    leaf_node_params.emplace_back(node->Min_x());
                    leaf_node_params.emplace_back(node->Min_y());
                    leaf_node_params.emplace_back(node->Size());
                    leaf_node_num++;
                }
                else{
                    if(node->ru_ != nullptr) q.push(node->ru_);
                    if(node->lu_ != nullptr) q.push(node->lu_);
                    if(node->rd_ != nullptr) q.push(node->rd_);
                    if(node->ld_ != nullptr) q.push(node->ld_);
                }
            }
        }

        map_file.write((char *)&(leaf_node_num), sizeof(int));

        ///如果地图尺寸较小，就以uint16来保存节点参数，以节省硬盘空间
        if(map_size_>65535){
            for(int i=0; i<leaf_node_num; i++){
                map_file.write((char *)&(leaf_node_params[3*i]), sizeof(int));
                map_file.write((char *)&(leaf_node_params[3*i+1]), sizeof(int));
                map_file.write((char *)&(leaf_node_params[3*i+2]), sizeof(int));
            }
        }
        else{
            for(int i=0; i<leaf_node_num; i++){
                uint16_t leaf_x_u16 = (uint16_t)leaf_node_params[3*i];
                uint16_t leaf_y_u16 = (uint16_t)leaf_node_params[3*i+1];
                uint16_t leaf_size_u16 = (uint16_t)leaf_node_params[3*i+2];
                map_file.write((char *)&(leaf_x_u16), sizeof(uint16_t));
                map_file.write((char *)&(leaf_y_u16), sizeof(uint16_t));
                map_file.write((char *)&(leaf_size_u16), sizeof(uint16_t));
            }
        }

        map_file.close();
        printf("Map saved at %s\n", map_file_path.c_str());
        return true;
    }
    else{
        printf("Unsupported save type. Supported type: PNG, PGM, BIN\n");
        return false;
    }
}

bool QuadtreeMap::LoadMap(const std::string& load_path,  const std::string& load_type){
    printf("Loading saved map file while clear all content of map besides layer ID\n");
    ///加载前先把现有地图内容清空
    Clear();
    ///截取文件名
    int p = -1;
    for(int i=load_path.size()-1;i>=0;i--) if(load_path[i]=='/') p=i+1;
    p = p==-1 ? 0:p;
    std::string file_name = load_path.substr(p);

    char type_name[3];
    sscanf(file_name.c_str(), "quadtreemap_%lf_%lf_%lf_%lf.%s", &center_x_, &center_y_, &center_yaw_, &grid_dim_, type_name);

    if(load_type=="PNG" || load_type=="PGM"){
        printf("Loading map file at %s... ...\n", load_path.c_str());

        cv::Mat m = cv::imread(load_path, cv::IMREAD_UNCHANGED);
        if(m.rows != m.cols) {
            std::cerr<<"Quadtree map image must be a square"<<std::endl;
            return false;
        }
        if(m.rows%2 != 0) {
            std::cerr<<"Quadtree map image side length must be even"<<std::endl;
            return false;
        }
        ///如果map_size不是2的整数次幂，就应该将其扩大
        int min_size = m.rows;
        max_depth_ = 0; map_size_ = 1;
        while (map_size_<min_size){
            map_size_ *= 2;
            max_depth_ += 1;
        }
        root_ = std::make_shared<QuadtreeNode>(0, 0, map_size_);  ///不要忘了重建root节点
        int offset = (map_size_ - min_size)/2; ///由于地图扩大造成的像素位置偏移

        for(int row = 0; row < map_size_; row++){
            uchar* ptr= m.ptr(row);
            for(int col = 0; col < map_size_; col++){
                if((int)ptr[col]==255){
                    AddLeafNode(map_size_-row-1+offset, col+offset, 1);
                }
            }
        }

        printf("Map file loaded\n");
        return true;
    }
    else if(load_type=="BIN"){
        printf("Loading map file at %s... ...\n", load_path.c_str());
        std::ifstream map_file;
        map_file.open(load_path, std::ios::binary);
        map_file.read((char *)&(map_size_), sizeof(int));
        root_ = std::make_shared<QuadtreeNode>(0, 0, map_size_); ///不要忘了重建root节点
        int leaf_node_num;
        map_file.read((char *)&(leaf_node_num), sizeof(int));
        ///从正层序排列的叶子节点恢复四叉树
        if(map_size_>65535){
            int leaf_node_min_x, leaf_node_min_y, leaf_node_size;
            for(int i=0; i<leaf_node_num; i++){
                map_file.read((char *)&(leaf_node_min_x), sizeof(int));
                map_file.read((char *)&(leaf_node_min_y), sizeof(int));
                map_file.read((char *)&(leaf_node_size), sizeof(int));
                AddLeafNode(leaf_node_min_x, leaf_node_min_y, leaf_node_size);
            }
        }
        else{
            uint16_t leaf_x_u16, leaf_y_u16, leaf_size_u16;
            for(int i=0; i<leaf_node_num; i++){
                map_file.read((char *)&(leaf_x_u16), sizeof(uint16_t));
                map_file.read((char *)&(leaf_y_u16), sizeof(uint16_t));
                map_file.read((char *)&(leaf_size_u16), sizeof(uint16_t));
                AddLeafNode((int)leaf_x_u16, (int)leaf_y_u16, (int)leaf_size_u16);
            }
        }

        map_file.close();
        printf("Map file loaded\n");
        return true;
    }
}

void QuadtreeMap::Clear(){
    root_.reset();
}

bool QuadtreeMap::toCVMat(cv::Mat &m){
    ///层序遍历，画大方块
    m = cv::Mat::zeros(map_size_, map_size_,CV_8UC1);
    std::queue< std::shared_ptr<QuadtreeNode> > q;
    q.push(root_);

    while (!q.empty()){
        int curr_level_size = q.size();
        for(int i=0; i<curr_level_size; i++){
            auto node = q.front(); q.pop();

            ///根据node的位置和大小画方块
            if( node->isLeaf() ){
                ///找到node对应的像素区域
                int start_row = map_size_ - node->Min_x() - node->Size();
                int start_col = node->Min_y();
                ///遍历像素
                for(int row = start_row; row < start_row+node->Size(); row++){
                    uchar* ptr= m.ptr(row);
                    for(int col = start_col; col < start_col+node->Size(); col++){
                        ptr[col]=255;
                    }
                }
            }
            else{
                if(node->ru_ != nullptr) q.push(node->ru_);
                if(node->lu_ != nullptr) q.push(node->lu_);
                if(node->rd_ != nullptr) q.push(node->rd_);
                if(node->ld_ != nullptr) q.push(node->ld_);
            }
        }
    }
    return true;
}

double QuadtreeMap::Center_x(){
    return center_x_;
}
double QuadtreeMap::Center_y(){
    return center_y_;
}
double QuadtreeMap::Center_yaw(){
    return center_yaw_;
}
double QuadtreeMap::Dim_x(){
    return map_size_*grid_dim_;
}
double QuadtreeMap::Dim_y(){
    return map_size_*grid_dim_;
}
int QuadtreeMap::LayerID(){
    return layer_id_;
}

bool QuadtreeMap::AddLeafNode(const int& min_x, const int& min_y, const int& size){
    std::shared_ptr< QuadtreeNode > curr_node = root_;
    int curr_depth = 0;
    int node_level = floor(log(map_size_/size)/log(2) );

    if(min_x<0 || min_y<0 || min_x+size>map_size_ || min_y+size>map_size_){
//        printf("Leaf node out of map range");
        return false;
    }
    ///找到叶子节点的父节点，如果没有就建立
    while(curr_depth < node_level){
        if( curr_node->isLeaf() ) return false; ///这里已经有叶子节点了，就不能添加了

        int center_row = curr_node->Min_x()+curr_node->Size()/2;
        int center_col = curr_node->Min_y()+curr_node->Size()/2;
        ///right up
        if(min_x >= center_row && min_y >= center_col){
            if(curr_node->ru_ == nullptr)
                curr_node->ru_ = std::make_shared< QuadtreeNode >
                        (curr_node, center_row, center_col,
                         curr_node->Size()/2);
            curr_node = curr_node->ru_;
            curr_depth += 1;
        }
        ///left up
        else if(min_x >= center_row && min_y < center_col){
            if(curr_node->lu_ == nullptr)
                curr_node->lu_ = std::make_shared< QuadtreeNode >
                        (curr_node, center_row, curr_node->Min_y(),
                         curr_node->Size()/2);
            curr_node = curr_node->lu_;
            curr_depth += 1;
        }
        ///right down
        else if(min_x < center_row && min_y >= center_col){
            if(curr_node->rd_ == nullptr)
                curr_node->rd_ = std::make_shared< QuadtreeNode >
                        (curr_node, curr_node->Min_x(), center_col,
                         curr_node->Size()/2);
            curr_node = curr_node->rd_;
            curr_depth += 1;
        }
        ///left down
        else{
            if(curr_node->ld_ == nullptr)
                curr_node->ld_ = std::make_shared< QuadtreeNode >
                        (curr_node, curr_node->Min_x(), curr_node->Min_y(),
                         curr_node->Size()/2);
            curr_node = curr_node->ld_;
            curr_depth += 1;
        }
    }
    ///此时curr_node已经指向对应的叶子节点
    curr_node->SetLeaf();

    ///合并节点
    while (curr_node->father_.lock() != nullptr){
        ///向上走一层
        curr_node = curr_node->father_.lock();
        if(curr_node->ru_!=nullptr && curr_node->lu_!=nullptr && curr_node->rd_!=nullptr && curr_node->ld_!=nullptr){
            if(curr_node->ru_->isLeaf() && curr_node->lu_->isLeaf() && curr_node->rd_->isLeaf() && curr_node->ld_->isLeaf()){
                curr_node->SetLeaf();
//                printf("merge node x:%d y:%d size:%d",curr_node->Min_x(), curr_node->Min_y(), curr_node->Size());
            }
        }
        else break;
    }
    return true;
}

